package svd;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;

public class SVDTrain {
	public static final int queryIdSize = 26243606;
	public static final int titleIdSize = 4051440;
	public static final int userIdSize = 23907635;
	public static final String pathFolder = "/home/ec2-user/yijia/kdd2012/version2/data/train/";
	public static final String testPath = "test";
	public static final int featureSize = 50;
	
	public static void trainNormalRSVD(String mode){
		int size1 = 0;
		int size2 = 0;
		String path = pathFolder + mode;
		if (mode.equals("q_t")){
			size1 = queryIdSize;
			size2 = titleIdSize;
		}else if (mode.equals("q_u")){
			size1 = queryIdSize;
			size2 = userIdSize;
		}else {
			size1 = userIdSize;
			size2 = titleIdSize;
		}
		int k = featureSize;
		double lrate = 0.001;
		double r = 0.0015;
		try{
			System.out.println("Read 1:");
			BufferedReader br = new BufferedReader(new FileReader(path));
			String line;
			int lineNum = 0;
			double sumI = 0;
			while ((line = br.readLine()) != null){
				if (lineNum % 1000000 == 0)
					System.out.println(lineNum);
				lineNum++;
				String[] paras = line.split("\t");
				int click = Integer.parseInt(paras[0]);
				int imp = Integer.parseInt(paras[1]);
				sumI += imp;
			}
			br.close();
			
			System.out.println("Read 2:");
			br = new BufferedReader(new FileReader(path));
			double[] ctr = new double[lineNum];
			int[] id1s = new int[lineNum];
			int[] id2s = new int[lineNum];
			
			double mean = sumI/lineNum;
			System.out.println("Mean :" + mean);
			lineNum = 0;
			while ((line = br.readLine()) != null){
				if (lineNum % 1000000 == 0)
					System.out.println(lineNum);
				String[] paras = line.split("\t");
				int click = Integer.parseInt(paras[0]);
				int imp = Integer.parseInt(paras[1]);
				ctr[lineNum] = (click*1.0 + 0.034882*mean) / (imp*1.0 + mean);
				
				int index = paras[2].indexOf("_");
				int id1 = Integer.parseInt(paras[2].substring(0, index));
				int id2 = Integer.parseInt(paras[2].substring(index+1));
				id1s[lineNum] = id1;
				id2s[lineNum] = id2;
				
				lineNum++;
			}
			br.close();
			
			//begin train
			System.out.println("begin training");
			double[][] f1 = new double[size1][k];
			double[][] f2 = new double[size2][k];
			for (int i = 0; i < size1; i++)
				for (int j = 0; j < k; j++)
					f1[i][j] = 0.1;
			
			for (int i = 0; i < size2; i++)
				for (int j = 0; j < k; j++)
					f2[i][j] = 0.1;
			
			double[] cache = new double[lineNum];
			for (int i = 0; i < k; i++){
				double lastrmse = 0;
				double rmse = 0;
				for (int iter = 0; iter < 100 || ((iter < 200 && lastrmse > rmse + 0.0001)); iter++){
					lastrmse = rmse;
					rmse = 0;
					for (int j = 0; j < lineNum; j++){
						if (j % 10000000 == 0)
							System.out.println(j + "\t" + rmse);
						int id1 = id1s[j];
						int id2 = id2s[j];
						
						double y = calcY(f1[id1][i], f2[id2][i], cache[j], i, true);
						double err = ctr[j] - y;
						
						f1[id1][i] += lrate*(err*f2[id2][i] - r*f1[id1][i]);
						f2[id2][i] += lrate*(err*f1[id1][i] - r*f2[id2][i]);
						
						rmse += err*err;
					}
					rmse = Math.sqrt(rmse/lineNum);
					System.out.println("Train " + i + " iter " + iter + 
							" rmse: " + rmse);
				}
				//update cache
				for (int j = 0; j < lineNum; j++){
					int id1 = id1s[j];
					int id2 = id2s[j];
					cache[j] += f1[id1][i]*f2[id2][i];
				}
			}
			
			//begin test
			System.out.println("SVD begin test...");
			br = new BufferedReader(new FileReader(testPath));
			BufferedWriter bw = new BufferedWriter(new FileWriter("svd_result"));
			int index = 0;
			while ((line = br.readLine()) != null){
				if (index % 1000000 == 0)
					System.out.println(index);
				index++;
				String[] paras = line.split("\t");
				int queryID = Integer.parseInt(paras[5]);
				int titleID = Integer.parseInt(paras[7]);
				int userID = Integer.parseInt(paras[9]);
				if (mode.equals("q_t")){
					if (queryID >= queryIdSize || titleID >= titleIdSize){
						bw.write("0.034882");
						bw.write("\n");
					}else{
						bw.write(String.valueOf(calcY(f1, f2, k, queryID, titleID)));
						bw.write("\n");
					}
				}else if (mode.equals("q_u")){
					if (queryID >= queryIdSize || userID >= userIdSize){
						bw.write("0.034882");
						bw.write("\n");
					}else{
						bw.write(String.valueOf(calcY(f1, f2, k, queryID, userID)));
						bw.write("\n");
					}
				}else {
					if (titleID >= titleIdSize || userID >= userIdSize){
						bw.write("0.034882");
						bw.write("\n");
					}else{
						bw.write(String.valueOf(calcY(f1, f2, k, userID, titleID)));
						bw.write("\n");
					}
				}
			}
			bw.flush();
			bw.close();
			
			//output to combine:
			System.out.println("SVD begin train_output...");
			br = new BufferedReader(new FileReader("training.txt"));
			bw = new BufferedWriter(new FileWriter("svd_result" + mode));
			index = 0;
			while ((line = br.readLine()) != null){
				if (index % 1000000 == 0)
					System.out.println(index);
				index++;
				String[] paras = line.split("\t");
				int queryID = Integer.parseInt(paras[7]);
				int titleID = Integer.parseInt(paras[9]);
				int userID = Integer.parseInt(paras[11]);
				
				if (mode.equals("q_t")){
					if (queryID >= queryIdSize || titleID >= titleIdSize){
						bw.write("0.034882");
						bw.write("\n");
					}else{
						bw.write(String.valueOf(calcY(f1, f2, k, queryID, titleID)));
						bw.write("\n");
					}
				}else if (mode.equals("q_u")){
					if (queryID >= queryIdSize || userID >= userIdSize){
						bw.write("0.034882");
						bw.write("\n");
					}else{
						bw.write(String.valueOf(calcY(f1, f2, k, queryID, userID)));
						bw.write("\n");
					}
				}else {
					if (titleID >= titleIdSize || userID >= userIdSize){
						bw.write("0.034882");
						bw.write("\n");
					}else{
						bw.write(String.valueOf(calcY(f1, f2, k, userID, titleID)));
						bw.write("\n");
					}
				}
			}
			bw.close();
		}catch(Exception e){
			e.printStackTrace();
		}
	}
	
	private static double calcY(double[][] f1, double[][] f2, int k, int id1, int id2){
		double sum = 0;
		for (int i = 0; i < k; i++)
			sum += f1[id1][i]*f2[id2][i];
		if (sum > 1) sum = 1;
		if (sum < 0) sum = 0;
		return sum;
	}
	
	private static double calcY(double d1, double d2, double c, int index, boolean isTraining) {
		double sum = c;
		sum += d1*d2;
		
		if (isTraining){
			sum += (featureSize - index - 1)*0.1*0.1;
		}
		
		if (sum < 0) sum = 0;
		if (sum > 1) sum = 1;
		return sum;
	}

	public static void trainThreeRSVD(String path1, String path2, String path3,
			int size1, int size2, int size3){
		
	}
	
	public static void trainBSVD(){
		
	}
	
	public static void trainThreeBSVD(){
		double lrate = 0.001;
		double r = 0.0015;
		String path = pathFolder + "qut";
		try{
			System.out.println("Read 1:");
			BufferedReader br = new BufferedReader(new FileReader(path));
			String line;
			int lineNum = 0;
			double sumI = 0;
			while ((line = br.readLine()) != null){
				if (lineNum % 1000000 == 0)
					System.out.println(lineNum);
				lineNum++;
				String[] paras = line.split("\t");
				int click = Integer.parseInt(paras[0]);
				int imp = Integer.parseInt(paras[1]);
				sumI += imp;
			}
			br.close();
			
			System.out.println("Read 2:");
			br = new BufferedReader(new FileReader(path));
			double[] ctr = new double[lineNum];
			int[] id1s = new int[lineNum];
			int[] id2s = new int[lineNum];
			int[] id3s = new int[lineNum];
			
			double mean = sumI/lineNum;
			System.out.println("Mean :" + mean);
			lineNum = 0;
			while ((line = br.readLine()) != null){
				if (lineNum % 1000000 == 0)
					System.out.println(lineNum);
				String[] paras = line.split("\t");
				int click = Integer.parseInt(paras[0]);
				int imp = Integer.parseInt(paras[1]);
				ctr[lineNum] = (click*1.0 + 0.034882*mean) / (imp*1.0 + mean);
				
				String[] segs = paras[2].split("_");
				id1s[lineNum] = Integer.parseInt(segs[0]);
				id2s[lineNum] = Integer.parseInt(segs[1]);
				id3s[lineNum] = Integer.parseInt(segs[2]);
				
				lineNum++;
			}
			br.close();
			
			//begin train
			System.out.println("begin training");
			double[][] fq = new double[queryIdSize][featureSize];
			double[][] fu = new double[userIdSize][featureSize];
			double[][] ft = new double[titleIdSize][featureSize];
			
			double[] bq = new double[queryIdSize];
			double[] bu = new double[userIdSize];
			double[] bt = new double[titleIdSize];
			
			for (int i = 0; i < queryIdSize; i++)
				for (int j = 0; j < featureSize; j++)
					fq[i][j] = 0.1;
			
			for (int i = 0; i < userIdSize; i++)
				for (int j = 0; j < featureSize; j++)
					fu[i][j] = 0.1;
			
			for (int i = 0; i < titleIdSize; i++)
				for (int j = 0; j < featureSize; j++)
					ft[i][j] = 0.1;
			
			double[] cache_qu = new double[lineNum];
			double[] cache_ut = new double[lineNum];
			double[] cache_qt = new double[lineNum];
			
			for (int i = 0; i < featureSize; i++){
				double lastrmse = 0;
				double rmse = 0;
				for (int iter = 0; iter < 100 || ((iter < 200 && lastrmse > rmse + 0.0001)); iter++){
					lastrmse = rmse;
					rmse = 0;
					for (int j = 0; j < lineNum; j++){
//						if (j % 10000000 == 0)
//							System.out.println(j + "\t" + rmse);
						int qid = id1s[j];
						int uid = id2s[j];
						int tid = id3s[j];
						
						double y = 0;
						y = y + 3*(featureSize - i - 1)*0.1*0.1;
						y = y + cache_qt[j] + cache_qu[j] + cache_ut[j];
						y = y + bq[qid] + bu[uid] + bt[tid];
						y += fq[qid][i]*ft[tid][i];
						y += fq[qid][i]*fu[uid][i];
						y += fu[uid][i]*ft[uid][i];
						
						if (y > 1) y = 1;
						if (y < 0) y = 0;
						
						double err = ctr[j] - y;
						
						bq[qid] += lrate*(err - r * bq[qid]);
						bt[tid] += lrate*(err - r * bt[tid]);
						bu[uid] += lrate*(err - r * bu[uid]);
						
						fq[qid][i] += lrate*(err*(fu[uid][i] + ft[tid][i]) - r*fq[qid][i]);
						ft[tid][i] += lrate*(err*(fu[uid][i] + fq[qid][i]) - r*ft[tid][i]);
						fu[uid][i] += lrate*(err*(fq[qid][i] + ft[tid][i]) - r*fu[uid][i]);
						
						rmse += err*err;
					}
					rmse = Math.sqrt(rmse/lineNum);
					System.out.println("Train " + i + " iter " + iter + 
							" rmse: " + rmse);
				}
				//update cache
				for (int j = 0; j < lineNum; j++){
					int qid = id1s[j];
					int uid = id2s[j];
					int tid = id3s[j];
					cache_qt[j] += fq[qid][i]*ft[tid][i];
					cache_qu[j] += fq[qid][i]*fu[uid][i];
					cache_ut[j] += fu[uid][i]*ft[tid][i];
				}
			}
			
			//begin test
			System.out.println("SVD begin test...");
			br = new BufferedReader(new FileReader(testPath));
			BufferedWriter bw = new BufferedWriter(new FileWriter("svd_result_qut"));
			int index = 0;
			while ((line = br.readLine()) != null){
				if (index % 1000000 == 0)
					System.out.println(index);
				index++;
				String[] paras = line.split("\t");
				int queryID = Integer.parseInt(paras[5]);
				int titleID = Integer.parseInt(paras[7]);
				int userID = Integer.parseInt(paras[9]);
				if (titleID >= titleIdSize || userID >= userIdSize || queryID >= queryIdSize) {
					bw.write("0.034882");
					bw.write("\n");
				} else {
					double y = 0.0;
					y = y + bq[queryID] + bt[titleID] + bu[userID];
					for (int i = 0; i < featureSize; i++){
						y += fq[queryID][i]*fu[userID][i];
						y += fq[queryID][i]*ft[titleID][i];
						y += fu[userID][i]*ft[titleID][i];
					}
					if (y > 1) y = 1.0;
					if (y < 0) y = 0.0;
					bw.write(String.format("%.6f", y));
					bw.write("\n");
				}
			}
			bw.flush();
			bw.close();
			
			//output to combine:
			System.out.println("SVD begin train_output...");
			br = new BufferedReader(new FileReader("training.txt"));
			bw = new BufferedWriter(new FileWriter("svd_result_train_qut"));
			index = 0;
			while ((line = br.readLine()) != null){
				if (index % 1000000 == 0)
					System.out.println(index);
				index++;
				String[] paras = line.split("\t");
				int queryID = Integer.parseInt(paras[7]);
				int titleID = Integer.parseInt(paras[9]);
				int userID = Integer.parseInt(paras[11]);
				
				if (titleID >= titleIdSize || userID >= userIdSize || queryID >= queryIdSize) {
					bw.write("0.034882");
					bw.write("\n");
				} else {
					double y = 0.0;
					y = y + bq[queryID] + bt[titleID] + bu[userID];
					for (int i = 0; i < featureSize; i++){
						y += fq[queryID][i]*fu[userID][i];
						y += fq[queryID][i]*ft[titleID][i];
						y += fu[userID][i]*ft[titleID][i];
					}
					if (y > 1) y = 1.0;
					if (y < 0) y = 0.0;
					bw.write(String.format("%.6f", y));
					bw.write("\n");
				}
				
			}
			bw.close();
			
		}catch(Exception e){
			e.printStackTrace();
		}
	}
	
	public static void main(String[] args){
		trainNormalRSVD("q_t");
	}

}
